/*******************************************************************************
#                                                                              #
#      MJPG-streamer allows to stream JPG frames from an input-plugin          #
#      to several output plugins                                               #
#                                                                              #
#      Copyright (C) 2007 Tom Stöveken                                         #
#                                                                              #
# This program is free software; you can redistribute it and/or modify         #
# it under the terms of the GNU General Public License as published by         #
# the Free Software Foundation; version 2 of the License.                      #
#                                                                              #
# This program is distributed in the hope that it will be useful,              #
# but WITHOUT ANY WARRANTY; without even the implied warranty of               #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the                #
# GNU General Public License for more details.                                 #
#                                                                              #
# You should have received a copy of the GNU General Public License            #
# along with this program; if not, write to the Free Software                  #
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA    #
#                                                                              #
*******************************************************************************/

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/ioctl.h>
#include <errno.h>
#include <signal.h>
#include <sys/socket.h>
#include <resolv.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <getopt.h>
#include <pthread.h>
#include <fcntl.h>
#include <time.h>
#include <syslog.h>
#include <dirent.h>

#include "../../utils.h"
#include "../../mjpg_streamer.h"

#define OUTPUT_PLUGIN_NAME "UDP output plugin"
#define STREAM_MAGIC "MJPG"  // 用于标识MJPG流的魔术字
#define MAX_CLIENTS 5        // 最大支持的客户端数量

static pthread_t worker;
static globals *pglobal;
static int fd, delay, max_frame_size;
static char *folder = "./tu";
static unsigned char *frame = NULL;
static char *command = NULL;
static int input_number = 0;
static int port = 0;

// 客户端信息结构体，用于保存请求视频流的客户端
typedef struct {
    struct sockaddr_in addr;
    socklen_t addr_len;
    int active;
    time_t last_seen;
} ClientInfo;

static ClientInfo clients[MAX_CLIENTS];
static pthread_mutex_t clients_mutex = PTHREAD_MUTEX_INITIALIZER;

/******************************************************************************
Description.: print a help message
******************************************************************************/
void help(void)
{
    fprintf(stderr, " ---------------------------------------------------------------\n" \
            " Help for output plugin..: "OUTPUT_PLUGIN_NAME"\n" \
            " ---------------------------------------------------------------\n" \
            " The following parameters can be passed to this plugin:\n\n" \
            " [-f | --folder ]........: folder to save pictures\n" \
            " [-d | --delay ].........: delay after saving pictures in ms\n" \
            " [-c | --command ].......: execute command after saveing picture\n" \
            " [-p | --port ]..........: UDP port to listen for commands\n" \
            "                           Supported commands:\n" \
            "                           - get_stream: Start receiving video stream\n" \
            "                           - get_photo [filename]: Capture current frame\n" \
            " [-i | --input ].........: read frames from the specified input plugin\n" \
            " ---------------------------------------------------------------\n");
}

/******************************************************************************
Description.: clean up allocated resources
******************************************************************************/
void worker_cleanup(void *arg)
{
    static unsigned char first_run = 1;

    if(!first_run) {
        DBG("already cleaned up resources\n");
        return;
    }

    first_run = 0;
    OPRINT("cleaning up resources allocated by worker thread\n");

    if(frame != NULL) {
        free(frame);
    }
    close(fd);
}

/******************************************************************************
Description.: 添加客户端到客户端列表
******************************************************************************/
void add_client(struct sockaddr_in *addr, socklen_t addr_len)
{
    pthread_mutex_lock(&clients_mutex);
    
    // 查找空位置或超时客户端
    for (int i = 0; i < MAX_CLIENTS; i++) {
        time_t now = time(NULL);
        // 客户端超时（30秒无活动）或未使用
        if (!clients[i].active || (now - clients[i].last_seen > 30)) {
            clients[i].addr = *addr;
            clients[i].addr_len = addr_len;
            clients[i].active = 1;
            clients[i].last_seen = now;
            DBG("Added client %d: %s:%d\n", i, inet_ntoa(addr->sin_addr), ntohs(addr->sin_port));
            break;
        }
    }
    
    pthread_mutex_unlock(&clients_mutex);
}

/******************************************************************************
Description.: 从客户端列表移除客户端
******************************************************************************/
void remove_client(struct sockaddr_in *addr)
{
    pthread_mutex_lock(&clients_mutex);
    
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].active && 
            (clients[i].addr.sin_addr.s_addr == addr->sin_addr.s_addr) &&
            (clients[i].addr.sin_port == addr->sin_port)) {
            clients[i].active = 0;
            DBG("Removed client %d: %s:%d\n", i, inet_ntoa(addr->sin_addr), ntohs(addr->sin_port));
            break;
        }
    }
    
    pthread_mutex_unlock(&clients_mutex);
}

/******************************************************************************
Description.: 向所有活跃客户端发送视频帧
******************************************************************************/
void send_frame_to_clients(int sd, unsigned char *frame, int frame_size)
{
    pthread_mutex_lock(&clients_mutex);
    
    // 构建带头部的数据包：魔术字(4字节) + 帧长度(4字节) + 帧数据
    unsigned char *packet = malloc(frame_size + 8);
    if (!packet) {
        DBG("Memory allocation failed for packet\n");
        pthread_mutex_unlock(&clients_mutex);
        return;
    }
    
    // 添加魔术字和帧长度
    memcpy(packet, STREAM_MAGIC, 4);
    uint32_t len = htonl(frame_size);  // 网络字节序
    memcpy(packet + 4, &len, 4);
    memcpy(packet + 8, frame, frame_size);
    
    time_t now = time(NULL);
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].active) {
            // 更新客户端活动时间
            clients[i].last_seen = now;
            
            // 发送帧数据
            if (sendto(sd, packet, frame_size + 8, 0, 
                      (struct sockaddr*)&clients[i].addr, 
                      clients[i].addr_len) < 0) {
                perror("sendto failed");
                // 发送失败，标记为不活跃
                clients[i].active = 0;
            } else {
                DBG("Sent frame to client %d (size: %d)\n", i, frame_size);
            }
        }
    }
    
    free(packet);
    pthread_mutex_unlock(&clients_mutex);
}

/******************************************************************************
Description.: 主工作线程
******************************************************************************/
void *worker_thread(void *arg)
{
    int ok = 1, frame_size = 0, rc = 0;
    char buffer1[1024] = {0};
    unsigned char *tmp_framebuffer = NULL;
    struct sockaddr_in addr;
    int sd;
    int bytes;
    unsigned int addr_len = sizeof(addr);
    char udpbuffer[1024] = {0};
    
    /* 设置清理函数 */
    pthread_cleanup_push(worker_cleanup, NULL);

    if(port <= 0) {
        OPRINT("必须提供有效的UDP端口\n");
        return NULL;
    }

    // 初始化客户端列表
    memset(clients, 0, sizeof(clients));

    // 创建UDP socket
    sd = socket(PF_INET, SOCK_DGRAM, 0);
    if (sd < 0) {
        perror("socket creation failed");
        return NULL;
    }
    
    bzero(&addr, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port = htons(8080);
    
    if(bind(sd, (struct sockaddr*)&addr, sizeof(addr)) != 0) {
        perror("bind failed");
        close(sd);
        return NULL;
    }

    // 设置socket为非阻塞模式，以便同时处理接收命令和发送视频流
    int flags = fcntl(sd, F_GETFL, 0);
    fcntl(sd, F_SETFL, flags | O_NONBLOCK);

    OPRINT("UDP输出插件已启动，端口: %d\n", port);
    OPRINT("支持的命令: get_stream, get_photo [文件名]\n");

    while(ok >= 0 && !pglobal->stop) {
        // 接收UDP命令
        memset(udpbuffer, 0, sizeof(udpbuffer));
        bytes = recvfrom(sd, udpbuffer, sizeof(udpbuffer)-1, 0, 
                        (struct sockaddr*)&addr, &addr_len);
        
        if (bytes > 0) {
            udpbuffer[bytes] = '\0';
            DBG("收到命令: %s 来自 %s:%d\n", udpbuffer, 
                inet_ntoa(addr.sin_addr), ntohs(addr.sin_port));

            // 处理get_stream命令 - 开始发送视频流
            if (strncmp(udpbuffer, "get_stream", 10) == 0) {
                add_client(&addr, addr_len);
                sendto(sd, "stream_started", 14, 0, 
                       (struct sockaddr*)&addr, addr_len);
            }
            // 处理stop_stream命令 - 停止发送视频流
            else if (strncmp(udpbuffer, "stop_stream", 11) == 0) {
                remove_client(&addr);
                sendto(sd, "stream_stopped", 14, 0, 
                       (struct sockaddr*)&addr, addr_len);
            }
            // 处理get_photo命令 - 捕获当前帧并保存
            else if (strncmp(udpbuffer, "get_photo", 9) == 0) {
                char *filename = strchr(udpbuffer, ' ');
                char fullpath[256];
                
                // 如果没有指定文件名，使用时间戳生成
                if (!filename || *(++filename) == '\0') {
                    time_t t = time(NULL);
                    struct tm *tm = localtime(&t);
                    snprintf(fullpath, sizeof(fullpath), 
                             "%s/photo_%04d%02d%02d_%02d%02d%02d.jpg",
                             folder, tm->tm_year+1900, tm->tm_mon+1, tm->tm_mday,
                             tm->tm_hour, tm->tm_min, tm->tm_sec);
                } else {
                    snprintf(fullpath, sizeof(fullpath), "%s/%s", folder, filename);
                }

                // 锁定并获取当前帧
                pthread_mutex_lock(&pglobal->in[input_number].db);
                pthread_cond_wait(&pglobal->in[input_number].db_update, 
                                 &pglobal->in[input_number].db);

                frame_size = pglobal->in[input_number].size;
                if(frame_size > max_frame_size) {
                    max_frame_size = frame_size + (1 << 16);
                    if((tmp_framebuffer = realloc(frame, max_frame_size)) == NULL) {
                        pthread_mutex_unlock(&pglobal->in[input_number].db);
                        LOG("内存不足\n");
                        return NULL;
                    }
                    frame = tmp_framebuffer;
                }
                memcpy(frame, pglobal->in[input_number].buf, frame_size);
                pthread_mutex_unlock(&pglobal->in[input_number].db);

                // 保存图片到文件
                if((fd = open(fullpath, O_CREAT | O_RDWR | O_TRUNC, 
                             S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH)) < 0) {
                    OPRINT("无法打开文件 %s\n", fullpath);
                    char errmsg[128];
                    snprintf(errmsg, sizeof(errmsg), "error: can't open file");
                    sendto(sd, errmsg, strlen(errmsg), 0, 
                           (struct sockaddr*)&addr, addr_len);
                } else {
                    if(write(fd, frame, frame_size) < 0) {
                        OPRINT("无法写入文件 %s\n", fullpath);
                        perror("write()");
                        char errmsg[128];
                        snprintf(errmsg, sizeof(errmsg), "error: write failed");
                        sendto(sd, errmsg, strlen(errmsg), 0, 
                               (struct sockaddr*)&addr, addr_len);
                    } else {
                        DBG("照片已保存: %s\n", fullpath);
                        sendto(sd, fullpath, strlen(fullpath), 0, 
                               (struct sockaddr*)&addr, addr_len);
                    }
                    close(fd);
                }

                // 执行用户指定的命令
                if(command != NULL) {
                    memset(buffer1, 0, sizeof(buffer1));
                    snprintf(buffer1, sizeof(buffer1), "%s \"%s\"", command, fullpath);
                    DBG("执行命令: %s", buffer1);

                    if((rc = setenv("MJPG_FILE", fullpath, 1)) != 0) {
                        LOG("setenv失败 (返回值 %d)\n", rc);
                    }

                    if((rc = system(buffer1)) != 0) {
                        LOG("命令执行失败 (返回值 %d)\n", rc);
                    }
                }

                if(delay > 0) {
                    usleep(1000 * delay);
                }
            }
            // 未知命令
            else {
                sendto(sd, "unknown_command", 15, 0, 
                       (struct sockaddr*)&addr, addr_len);
            }
        }

        // 获取最新帧并发送给所有活跃客户端
        pthread_mutex_lock(&pglobal->in[input_number].db);
        // 使用timedwait避免一直阻塞，以便能响应stop信号
        struct timespec timeout;
        clock_gettime(CLOCK_REALTIME, &timeout);
        timeout.tv_sec += 1;  // 1秒超时
        
        if (pthread_cond_timedwait(&pglobal->in[input_number].db_update, 
                                  &pglobal->in[input_number].db, &timeout) == 0) {
            // 成功获取新帧
            frame_size = pglobal->in[input_number].size;

            // 确保缓冲区足够大
            if(frame_size > max_frame_size) {
                DBG("增加缓冲区大小到 %d\n", frame_size);
                max_frame_size = frame_size + (1 << 16);
                if((tmp_framebuffer = realloc(frame, max_frame_size)) == NULL) {
                    pthread_mutex_unlock(&pglobal->in[input_number].db);
                    LOG("内存不足\n");
                    return NULL;
                }
                frame = tmp_framebuffer;
            }

            // 复制帧数据
            memcpy(frame, pglobal->in[input_number].buf, frame_size);
            
            // 发送帧到所有活跃客户端
            send_frame_to_clients(sd, frame, frame_size);
        }
        pthread_mutex_unlock(&pglobal->in[input_number].db);
    }

    // 清理资源
    close(sd);
    pthread_cleanup_pop(1);

    return NULL;
}

/*** 插件接口函数 ***/
/******************************************************************************
Description.: 初始化插件
******************************************************************************/
int output_init(output_parameter *param)
{
    int i;

    delay = 0;
    param->argv[0] = OUTPUT_PLUGIN_NAME;

    /* 显示所有参数用于调试 */
    for(i = 0; i < param->argc; i++) {
        DBG("argv[%d]=%s\n", i, param->argv[i]);
    }

    reset_getopt();
    while(1) {
        int option_index = 0, c = 0;
        static struct option long_options[] = {
            {"h", no_argument, 0, 0},
            {"help", no_argument, 0, 0},
            {"f", required_argument, 0, 0},
            {"folder", required_argument, 0, 0},
            {"d", required_argument, 0, 0},
            {"delay", required_argument, 0, 0},
            {"c", required_argument, 0, 0},
            {"command", required_argument, 0, 0},
            {"p", required_argument, 0, 0},
            {"port", required_argument, 0, 0},
            {"i", required_argument, 0, 0},
            {"input", required_argument, 0, 0},
            {0, 0, 0, 0}
        };

        c = getopt_long_only(param->argc, param->argv, "", long_options, &option_index);
        if(c == -1) break;

        if(c == '?') {
            help();
            return 1;
        }

        switch(option_index) {
            /* 帮助选项 */
        case 0:
        case 1:
            help();
            return 1;
            break;

            /* 文件夹选项 */
        case 2:
        case 3:
            folder = malloc(strlen(optarg) + 1);
            strcpy(folder, optarg);
            if(folder[strlen(folder)-1] == '/')
                folder[strlen(folder)-1] = '\0';
            break;

            /* 延迟选项 */
        case 4:
        case 5:
            delay = atoi(optarg);
            break;

            /* 命令选项 */
        case 6:
        case 7:
            command = strdup(optarg);
            break;

            /* 端口选项 */
        case 8:
        case 9:
            port = atoi(optarg);
            break;

            /* 输入插件选项 */
        case 10:
        case 11:
            input_number = atoi(optarg);
            break;
        }
    }

    pglobal = param->global;
    if(!(input_number < pglobal->incnt)) {
        OPRINT("错误: 输入插件编号 %d 超出范围，只加载了 %d 个插件\n", 
               input_number, pglobal->incnt);
        return 1;
    }
    
    OPRINT("输入插件.........: %d: %s\n", input_number, pglobal->in[input_number].plugin);
    OPRINT("输出文件夹.......: %s\n", folder);
    OPRINT("保存后延迟.......: %d ms\n", delay);
    OPRINT("执行命令.........: %s\n", (command == NULL) ? "禁用" : command);
    OPRINT("UDP端口..........: %d\n", port);
    
    return 0;
}

/******************************************************************************
Description.: 停止工作线程
******************************************************************************/
int output_stop(int id)
{
    DBG("取消工作线程\n");
    pthread_cancel(worker);
    return 0;
}

/******************************************************************************
Description.: 创建并启动工作线程
******************************************************************************/
int output_run(int id)
{
    DBG("启动工作线程\n");
    pthread_create(&worker, 0, worker_thread, NULL);
    pthread_detach(worker);
    return 0;
}
